function w = est_cv(y,X,TUN,option,CV)
% estimate function by cross validation
% TUN is a the tuning vector
%   option =    1, LASSO 
%             1.5, LASSO (without sum to 1)
%               2, RIDGE 
%             2.5, RIDGE (without sum to 1)
%               3, L2relax (no shrinkage); 
%               4, L2relax (linear shrinkage); 
%               5, L2relax (nonlinear shrinkage); 

    
    % 0. house clean
    if length(TUN) > 1
        if nargin < 5
            CV = 1; % default, iid
        end
        % 1. five fold cross validation
        cvMse = zeros(length(TUN),1);            
        if CV == 1
            for i = 1:length(TUN)
                regf=@(XTRAIN,ytrain,XTEST)(cvfun(XTRAIN,ytrain,XTEST,TUN(i),option));
                cvMse(i,1) = crossval('mse',X,y,'predfun',regf,'kfold',5);
            end    
        elseif CV == 2
            for i = 1:length(TUN)            
                cvMse(i,1) = nest_cv(y,X,TUN(i),option);
            end    
        end
        % 2. find the best
        [~,IN] = min(cvMse);
        tun = TUN(IN);
    else
        tun = TUN;
    end
    N = size(X,2);
    options = optimset('display','none');
    w0 = ones(N,1)/N;
   	if option == 1
        w = fmincon(@(w) lasso_fun(w,X,y,tun),w0,[],[],ones(1,N),1,[],[],[],options);
    elseif option == 1.5
        w = lasso(X,y,'Lambda',tun);
    elseif option == 2
        w = fmincon(@(w) ridge_fun(w,X,y,tun),w0,[],[],ones(1,N),1,[],[],[],options);
    elseif option == 2.5
        w = ridge(y,X,tun);
    elseif option == 3
        w = l2relax0(y,X,tun,1);    % no shrinkage
    elseif option == 4
        w = l2relax0(y,X,tun,2);    % linear shrinkage
    elseif option == 5
        w = l2relax0(y,X,tun,3);    % nonlinear shrinkage
    end
end

function yfit = cvfun(XTRAIN,ytrain,XTEST,tun,option)
% cross validation estimation function
    warning off
    N = size(XTRAIN,2);
    options = optimset('display','none');
    w0 = ones(N,1)/N;
    if option == 1 % lasso    
        w = fmincon(@(w) lasso_fun(w,XTRAIN,ytrain,tun),w0,[],[],ones(1,N),1,[],[],[],options);
        yfit = XTEST*w;
    elseif option == 1.5 % lasso (no sum to 1)
        w = lasso(XTRAIN,ytrain,'Lambda',tun);
        yfit = XTEST*w;
    elseif option == 2 % ridge
        w = fmincon(@(w) ridge_fun(w,XTRAIN,ytrain,tun),w0,[],[],ones(1,N),1,[],[],[],options);
        yfit = XTEST*w;
    elseif option == 2.5 % ridge (no sum to 1)        
        w = ridge(ytrain,XTRAIN,tun);    
        yfit = XTEST*w;
    elseif option == 3 % l2relax (no shrinkage)
        w = l2relax0(ytrain,XTRAIN,tun,1);
        yfit = XTEST*w;
    elseif option == 4 % l2relax (linear shrinkage)
        w = l2relax0(ytrain,XTRAIN,tun,2);
        yfit = XTEST*w;
    elseif option == 5 % l2relax (nonlinear shrinkage)
        w = l2relax0(ytrain,XTRAIN,tun,3);
        yfit = XTEST*w;
    end
end

function cri = lasso_fun(w,X,y,tun)
    N = size(X,2);
    E = y-X*w;    
    SIG = analytical_shrinkage(E);
    cri = SIG/2+tun*sum(abs(w-1/N));        % recentered around 1/N
end

function cri = ridge_fun(w,X,y,tun)
    N = size(X,2);
    E = y-X*w;    
    SIG = analytical_shrinkage(E);
    cri = SIG/2+tun*((w-1/N)'*(w-1/N));     % recentered around 1/N
end

function MSE = nest_cv(y,X,tun,option)
    % 1. construct 5 nested index
    n = length(y);
    N = size(X,2);
    options = optimset('display','none');
    w0 = ones(N,1)/N;
    d = round(n/5);
    tmp = 1:d:n;
    IN = [tmp(1:5)' [tmp(2:5)'-1; n]];
    INt = [ones(4,1) IN(1:4,2)];
    INe = IN(2:end,:);
    if d < 12         % to avoid the "analytical_shrinkage.m" error
        INt = INt(2:end,:);
        INe = INe(2:end,:);
    end
    % 2. obtain MSE for each tuning value
    for j = 1:size(INt,1)
        yt = y(INt(j,1):INt(j,2),:);
        Xt = X(INt(j,1):INt(j,2),:);
        ye = y(INe(j,1):INe(j,2),:);
        Xe = X(INe(j,1):INe(j,2),:);
        if option == 1
            wt = fmincon(@(wt) lasso_fun(wt,Xt,yt,tun),w0,[],[],ones(1,N),1,[],[],[],options);
        elseif option == 1.5
            wt = lasso(Xt,yt,'Lambda',tun);    
        elseif option == 2
            wt = fmincon(@(wt) ridge_fun(wt,Xt,yt,tun),w0,[],[],ones(1,N),1,[],[],[],options);
        elseif option == 2.5
            wt = ridge(yt,Xt,tun);    
        elseif option == 3
            wt = l2relax0(yt,Xt,tun,1);     % no shrinkage
        elseif option == 4
            wt = l2relax0(yt,Xt,tun,2);     % linear shrinkage
        elseif option == 5
            wt = l2relax0(yt,Xt,tun,3);     % nonlinear shrinkage
        elseif option == 6
            wt = l2relax2(yt,Xt,tun,3);     % nonlinear without summing to 1 restriction
        end
        MSEt(j) = mean((ye-Xe*wt).^2); 
    end
    MSE = mean(MSEt);    
end
 


